# -*- coding: utf-8 -*-
"""
   File Name:  test.py
   Author :    liccoo
   Time:       2022/8/24 12:48
"""
import argparse
import os

import torch

from utils.load_dataset import ForwardDataset, InverseDataset
from utils.path import logs_exp, extract_model_path
from utils.plot import plot_forward, plot_inverse, plot_inverse_scatter
from utils.print_save_hyperparameter import print_save_hyperparas

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # 通用超参数
    parser.add_argument('--plot-num', type=int, default=30, help='The total number of graphs to plot')
    parser.add_argument('--random-sample', action='store_false',
                        help='Value indicates whether the sample is random, defaule=True.')
    parser.add_argument('--exp-order', type=int, default=0,
                        help='The path of the model saved,default=0 indicates the latest exp.')
    parser.add_argument('--model-order', type=int, default=0,
                        help='The order of the trained models which are saved. default=0 indicates the latest model')
    parser.add_argument('--device', type=str, default='cpu',
                        help='cuda:0, cuda:2, ..., cuda:n, cpu.')
    parser.add_argument('--new-dir', type=str, default='',
                        help='Optional, the additional name of the parent directory path of exp')
    # 定制超参数
    parser.add_argument('--model', type=str, required=True, help='|forward|inverse|tandem|')
    args = parser.parse_args()

    # 检查 args 参数
    # check_args_test(args)

    # 获取当前项目的绝对路径
    project_dir_path = os.path.dirname(os.path.realpath('__file__'))

    # 测试集的路径
    test_dataset_path = os.path.join(project_dir_path, 'dataset', 'test_dataset.csv')

    # 设置存放测试 log 文件的根目录
    logs_exp_path = logs_exp(project_dir_path, args, 'test')

    # 加载网络模型路径
    model_path = extract_model_path(project_dir_path, args.model, args.exp_order, args.model_order)

    # 打印和保存超参数
    print_save_hyperparas(args, logs_exp_path)

    # 加载训练好的网络模型
    if args.device == 'cpu':
        model = torch.load(model_path, map_location=torch.device('cpu'))
    else:
        model = torch.load(model_path)

    # 对相应的模型画图
    if args.model == 'forward':
        # 数据集 dataset 类实例化
        dataset = ForwardDataset(test_dataset_path)
        plot_forward(args, model, dataset, logs_exp_path, 0, 1, 1)

    elif args.model in ['inverse', 'tandem']:
        # 数据集 dataset 类实例化
        dataset = InverseDataset(test_dataset_path)
        # plot_inverse(args, model, dataset, logs_exp_path, 0, 1, 1)
        plot_inverse_scatter(args, model, dataset, logs_exp_path)
